import json
from pathlib import Path
from typing import Optional
from tenacity import retry, stop_after_attempt, wait_random_exponential
from metagpt.utils.common import CodeParser
from .pylint_examinate import PylintRunner

from metagpt.actions import Action, ActionOutput
# from metagpt.actions.refine_an import (
#     REFINE_DESIGN_API_NODE
# )
from metagpt.actions.design_api_an import (
    REFINED_DESIGN_NODE
)
from metagpt.actions.project_management_an import (
    REFINED_PM_NODE
)
# from metagpt.const import DATA_API_DESIGN_FILE_REPO, SEQ_FLOW_FILE_REPO
from metagpt.logs import logger
from metagpt.schema import Document, Documents, Message

SYSTEM_DESIGN_REFINE_TEMPLATE="""
### PRD(Product Requirement Document)
{prd}
### System Design
{system_design}
### Refine Suggestions
{suggestions}
"""

TASK_LIST_REFINE_TEMPLATE="""
### System Design
{system_design}
### Task_List and Dependencies
{task_list}
### Refine Suggestions
{suggestions}
"""

ENGINEER_CODE_REFINE="""
NOTICE
Role: You are a professional engineer; The main goal is to complete incremental development by combining legacy code and plan and Incremental Change, ensuring the integration of new features.
### System Design
{system_design}

### Task_List and Dependencies
{task_list}

### Original Code Content
Code in {filename}:
{code_content}

### Refine Suggestions
{suggestions}

### Lint
{lint}

# Instruction: Based on the context, rewrite code.
## Rewrite Code: rewrite file {filename}, refine the code to ensure satisified the suggestions, system_design, task_list, and fix bugs according Lint
1. Code should fix all bugs in Lint.
2. Code should match the suggestions.
3. give the complete code that could be run.

"""


class Refine(Action):
    name: str = ""
    i_context: Optional[str] = None
    desc: str = (
        ""
    )

    async def run(self, refine_idx:int, suggestions: str, with_messages: Message, schema: str = None):
        operfunc = [
            self._system_design_refine,
            self._task_list_refine,
            self._engineer_code_refine,
        ]
        lenth = len(operfunc)
        if refine_idx < 0 or refine_idx >= lenth:
            raise Exception(f"error refine_idx: {refine_idx}, length: {lenth}")
        return await operfunc[refine_idx](suggestions, with_messages, schema)

    async def _system_design_refine(self, suggestions: str, with_messages: Message, schema: str = None):
        """
        system design refine
        """

        async def _new_refine(llm, prd_doc, system_design_doc,suggestions):
            context = SYSTEM_DESIGN_REFINE_TEMPLATE.format(
                prd=prd_doc.content,
                system_design=system_design_doc.content,
                suggestions=suggestions
            )
            refine = await REFINED_DESIGN_NODE.fill(context=context, llm=llm)
            return refine

        # Use `git status` to identify which PRD documents have been modified in the `docs/prd` directory.
        changed_prds = self.repo.docs.prd.changed_files
        # Use `git status` to identify which design documents in the `docs/system_designs` directory have undergone
        # changes.
        changed_system_designs = self.repo.docs.system_design.changed_files

        prd = None
        old_system_design_doc = None
        for filename in changed_system_designs.keys():
            prd = await self.repo.docs.prd.get(filename)
            old_system_design_doc = await self.repo.docs.system_design.get(filename)
            new_system_design_doc = await _new_refine(self.llm,prd,old_system_design_doc,suggestions)
            doc = await self.repo.docs.system_design.save(
                filename=filename,
                content=new_system_design_doc.instruct_content.model_dump_json(),
                dependencies={prd.root_relative_path},
            )
            logger.info(f"new system design: {doc}")
        return ActionOutput(content=doc.model_dump_json(),instruct_content=doc)

    async def _task_list_refine(self, suggestions: str, with_messages: Message, schema: str =None):

        async def _new_refine(llm, system_design_doc, task_list_doc,suggestions):
            context = TASK_LIST_REFINE_TEMPLATE.format(
                system_design=system_design_doc.content,
                task_list=task_list_doc.content,
                suggestions=suggestions
            )
            refine = await REFINED_PM_NODE.fill(context=context, llm=llm)
            return refine

        changed_system_designs = self.repo.docs.system_design.changed_files

        system_design = None
        task_list = None 

        for filename in changed_system_designs.keys():
            system_design = await self.repo.docs.system_design.get(filename)
            task_list = await self.repo.docs.task.get(filename)
            new_task_list_doc = await _new_refine(self.llm,system_design, task_list,suggestions)
            doc = await self.repo.docs.task.save(
                filename=filename,
                content=new_task_list_doc.instruct_content.model_dump_json(),
                dependencies={system_design.root_relative_path},
            )
            logger.info(f"new system design: {doc}")
        return ActionOutput(content=doc.model_dump_json(),instruct_content=doc)

    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
    async def write_code(self, prompt) -> str:
        code_rsp = await self._aask(prompt)
        code = CodeParser.parse_code(block="", text=code_rsp)
        return code

    async def _engineer_code_refine(self, suggestions: dict, with_messages: Message, schema: str=None):

        # suggestion parse
        suggestions = dict(suggestions)

        # build code content
        changed_code_files_repo = self.repo.with_src_path(self.context.src_workspace)
        src_file_repo = changed_code_files_repo.srcs
        changed_src_files = list(src_file_repo.all_files)

        runner = PylintRunner(str(src_file_repo.workdir.absolute()))
        lintdict = runner.run()

        async def _new_refine(
                llm,system_design_doc,task_list_doc,filename,code_content,suggestions,lint
            ):
            context = ENGINEER_CODE_REFINE.format(
                system_design=system_design_doc.content,
                task_list=task_list_doc.content,
                filename=filename,
                code_content=code_content,
                suggestions=suggestions,
                lint=lint
            )
            new_code_content = await self.write_code(context)
            # refine = await REFINED_ENGINEER_NODE.fill(context=context, llm=llm)
            TAGLINE = "## THIS IS THE REFINE CODE LINE ##"
            return TAGLINE + "\n" + new_code_content

        # Use `git status` to identify which PRD documents have been modified in the `docs/prd` directory.
        changed_task = self.repo.docs.task.changed_files
        # Use `git status` to identify which design documents in the `docs/system_designs` directory have undergone
        # changes.
        changed_system_designs = self.repo.docs.system_design.changed_files

        system_design_doc = None
        task_doc = None
        for filename in changed_system_designs.keys():
            system_design_doc = await self.repo.docs.system_design.get(filename)
            task_doc = await self.repo.docs.task.get(filename)
            break

        # 逐个代码修改
        for filename in changed_src_files:
            code_doc = await src_file_repo.get(filename=filename)
            # If the file does not exist in the src workspace, skip it
            if not code_doc:
                continue

            cur_code_suggestion = suggestions.get(filename,None)
            cur_lint = lintdict.get(filename,None)

            if cur_code_suggestion or cur_lint:
                new_code_content = await _new_refine(
                    self.llm,system_design_doc,task_doc,filename,code_doc.content,
                    cur_code_suggestion,cur_lint
                )

                doc = await self.repo.with_src_path(self.context.src_workspace).srcs.save(
                    filename=filename,
                    # content=new_code_doc.instruct_content.model_dump_json(),
                    content=new_code_content,
                    dependencies={task_doc.root_relative_path}
                )
                logger.info(f"refined code for {filename}:\n {new_code_content}")
                logger.info(f"refined doc {doc.filename}:\n {doc.content}")

        # return ActionOutput(content=No,instruct_content=doc)
        return None


